//
// Copyright (c) Microsoft Corporation.  All rights reserved.
//
//
// Use of this sample source code is subject to the terms of the Microsoft
// license agreement under which you licensed this sample source code. If
// you did not accept the terms of the license agreement, you are not
// authorized to use this sample source code. For the terms of the license,
// please see the license agreement between you and Microsoft or, if applicable,
// see the LICENSE.RTF on your install media or the root of your tools installation.
// THE SAMPLE SOURCE CODE IS PROVIDED "AS IS", WITH NO WARRANTIES OR INDEMNITIES.
//
#include <windows.h>
#include <pehdr.h>
#include <romLdr.h>
#include <oal.h>
#include <bootStore.h>
#include <bootStoreUtils.h>
#include <cecompress.h>

//------------------------------------------------------------------------------

#define BINCOMPRESS_SIGNATURE       "XPRS"
#define NO_COMPRESSION_BIT          (1 << 31)
#define PARTIAL_BLOCK_BIT           (1 << 30)

//------------------------------------------------------------------------------
// Data buffers

static
UINT8
s_readBuffer[CECOMPRESS_MAX_BLOCK];

static
BOOL
s_bufferFull = FALSE;

static
UINT8
s_allocBuffer[0x2000];

//------------------------------------------------------------------------------

static
VOID*
BootStoreCompressAlloc(
    VOID* pContext,
    DWORD size
    )
{
    VOID *pMemory = NULL;
        
    if (!s_bufferFull && (size <= sizeof(s_allocBuffer)))
        {
        pMemory = s_allocBuffer;
        s_bufferFull = TRUE;
        }
    return pMemory;
}

//------------------------------------------------------------------------------

static
VOID
BootStoreCompressFree(
    VOID* pContext, 
    VOID* pMemory
    )
{
    s_bufferFull = FALSE;
}

//------------------------------------------------------------------------------

UINT32
BootStoreReadBinCompressed(
    HANDLE hStore,
    VOID *pBuffer,
    UINT32 bufferSize
    )
{
    UINT32 rc = 0;
    CeCompressDecodeStream decodeStream = NULL;
    UINT8 sign[dimof(BINCOMPRESS_SIGNATURE) - 1];
    UINT32 totalLength;
    UINT32 header;
    UINT32 compressedLength, decompressedLength;
    BOOL storedBlock, partialBlock;
    

    // Initialize the decoder
    decodeStream = CeCompressDecodeCreate(NULL, BootStoreCompressAlloc); 
    if (decodeStream == NULL) goto cleanUp;
    
    // Read signature
    if (sizeof(sign) != BootStoreRead(hStore, sign, sizeof(sign)))
        {
        OALMSG(OAL_ERROR, (L"ERROR: BootStoreReadBinCompressed: "
            L"Unexpected stream read error!\r\n"
            ));
        goto cleanUp;
        }
    
    // Check signature value
    if (memcmp(sign, BINCOMPRESS_SIGNATURE, sizeof(sign)) != 0)
        {
        OALMSG(OAL_ERROR, (L"ERROR: BootStoreReadBinCompressed: "
            L"BinCompress signature XPRS doesn't found!\r\n"
            ));
        goto cleanUp;
        }

    // Read decompressed file length
    if (sizeof(totalLength) != BootStoreRead(
            hStore, &totalLength, sizeof(totalLength)))
        {
        OALMSG(OAL_ERROR, (L"ERROR: BootStoreReadBinCompressed: "
            L"Unexpected stream read error!"
            ));
        goto cleanUp;
        }
    
    // Decode input blocks into CECOMPRESS_MAX_BLOCK-sized blocks
    while (sizeof(header) == BootStoreRead(hStore, &header, sizeof(header)))
        {
        storedBlock = header & NO_COMPRESSION_BIT;
        partialBlock = header & PARTIAL_BLOCK_BIT;
        compressedLength = header & ~(NO_COMPRESSION_BIT|PARTIAL_BLOCK_BIT);

        // Read decompressed block length if it's a partial block
        if (partialBlock)
            {
            if (sizeof(decompressedLength) != BootStoreRead(
                    hStore, &decompressedLength, sizeof(decompressedLength)))
                {
                OALMSG(OAL_ERROR, (L"ERROR: BootStoreReadBinCompressed: "
                    L"Unexpected stream read error!\r\n"
                    ));
                goto cleanUp;
                }
            }
        else
            {
            decompressedLength = CECOMPRESS_MAX_BLOCK;
            }

        // Check if decompressed block fits in buffer
        if (decompressedLength > bufferSize) break;

        // Read compressed data
        if (compressedLength != BootStoreRead(
                hStore, s_readBuffer, compressedLength))
            {
            OALMSG(OAL_ERROR, (L"ERROR: BootStoreReadBinCompressed: "
                L"Unexpected boot store stream read error!\r\n"
                ));
            goto cleanUp;
            }

        // Compressed or just stored?
        if (storedBlock)
            {
            // Uncompressed block - just copy
            memcpy(pBuffer, s_readBuffer, compressedLength);
            }
        else
            {
            // Decode the data
            if (CeCompressDecode(
                    decodeStream, pBuffer, decompressedLength, 
                    decompressedLength, s_readBuffer, compressedLength
                    ) == -1)
                {
                OALMSG(OAL_ERROR, (L"ERROR: BootStoreReadBinCompressed: "
                    L"Decompression failed!\r\n"
                    ));
                goto cleanUp;
                }
            }

        // Move forward in buffer
        (UINT8*)pBuffer += decompressedLength;
        bufferSize -= decompressedLength;
        rc += decompressedLength;
        }
   
cleanUp:

    // Clean things up
    if (decodeStream != NULL) 
        {
        CeCompressDecodeClose(decodeStream, NULL, BootStoreCompressFree);
        }
    return rc;
}

//------------------------------------------------------------------------------

